package cn.sskxyz.mongodb.locks;

import java.time.Instant;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.mongodb.core.FindAndModifyOptions;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;

import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;

public class MongoDistributedLock implements DistributedLock {

	private static final Logger logger = LoggerFactory.getLogger(MongoDistributedLock.class);

	private MongoTemplate mongoTemplate;

	public MongoDistributedLock(MongoTemplate mongoTemplate) {
		this.mongoTemplate = mongoTemplate;
	}

	@Override
	public String acquire(String key, long expiration, TimeUnit timeUnit) {
		Query query = Query.query(Criteria.where("_id").is(key));
		String token = UUID.randomUUID().toString();
		long expireAt = Instant.now().plusMillis(timeUnit.toMillis(expiration)).toEpochMilli();
		Update update = new Update().setOnInsert("_id", key).setOnInsert("expireAt", expireAt)
				.setOnInsert("token", token);

		FindAndModifyOptions options = new FindAndModifyOptions().upsert(true).returnNew(true);
		LockDocument doc = mongoTemplate.findAndModify(query, update, options, LockDocument.class);
		boolean locked = doc.getToken() != null && doc.getToken().equals(token);

		if (!locked && doc.getExpireAt() < Instant.now().toEpochMilli()) {
			DeleteResult deleted = this.mongoTemplate.remove(Query.query(Criteria.where("_id").is(key).and("token")
					.is(doc.getToken()).and("expireAt").is(doc.getExpireAt())), LockDocument.class);
			if (deleted.getDeletedCount() >= 1) {
				return this.acquire(key, expiration, timeUnit);
			}
		}

		logger.debug("Tried to acquire lock for key {} with token {} . Locked: {}", key, token, locked);
		return locked ? token : null;
	}

	@Override
	public boolean release(String key, String token) {
		Query query = Query.query(Criteria.where("_id").is(key).and("token").is(token));
		DeleteResult deleted = mongoTemplate.remove(query, LockDocument.class);
		boolean released = deleted.getDeletedCount() == 1;
		if (released) {
			logger.debug("Remove query successfully affected 1 record for key {} with token {}", key, token);
		} else if (deleted.getDeletedCount() > 0) {
			logger.error("Unexpected result from release for key {} with token {}, released {}", key, token, deleted);
		} else {
			logger.error("Remove query did not affect any records for key {} with token {}", key, token);
		}

		return released;
	}

	@Override
	public boolean refresh(String key, String token, long expiration, TimeUnit timeUnit) {
		Query query = Query.query(Criteria.where("_id").is(key).and("token").is(token));
		long expireAt = Instant.now().plusMillis(timeUnit.toMillis(expiration)).toEpochMilli();
		Update update = Update.update("expireAt", expireAt);
		UpdateResult updated = mongoTemplate.updateFirst(query, update, LockDocument.class);

		final boolean refreshed = updated.getModifiedCount() == 1;
		if (refreshed) {
			logger.debug("Refresh query successfully affected 1 record for key {} " + "with token {}", key, token);
		} else if (updated.getModifiedCount() > 0) {
			logger.error("Unexpected result from refresh for key {} with token {}, " + "released {}", key, token, updated);
		} else {
			logger.warn("Refresh query did not affect any records for key {} with token {}. "
					+ "This is possible when refresh interval fires for the final time "
					+ "after the lock has been released", key, token);
		}
		return refreshed;
	}
}
